# -*- encoding: utf-8 -*-
'''

Main entry

@File    :   main.py
@Time    :   2021/12/04 16:20:24
@Author  :   Yangshuo He
@Contact :   sugarhe58@gmail.com
'''

import argparse
import cv2
import os
import time
import datetime
import picamera
from picamera.array import PiRGBArray
from picamera import PiCamera
from model import OpenvinoModel
from message import Message
from utils import myLog, checkPi

def parse_args() -> argparse.Namespace:
    """Parse and return command line arguments"""
    parser = argparse.ArgumentParser(add_help=False)
    args = parser.add_argument_group("Options")
    # fmt: off
    args.add_argument("-h", "--help", action="help", help="Show this help message and exit.")
    args.add_argument("-mn", "--model_name", type=str,
                        default="face-detection-0200",
                        help="Optional. Type of a trained model."
                            "Default value is 'person-detection-0200'" 
                            "Suitable for:"
                            "pedestrian-detection-adas-0002"
                            "person-detection-0200"
                            "person-detection-0201"
                            "person-detection-0202"
                            "face-detection-0200"
                            "face-detection-0202"
                            "face-detection-0204"
                            "face-detection-adas-0001"
                            "face-detection-retail-0004"
                            "face-detection-retail-0005"
                    )
    args.add_argument("-md","--model_dir", type=str, 
                        default="/home/pi/openvino/open_model_zoo/models/intel/", 
                        help="Optional. Path to the models directory"
                            "Default value is '/home/pi/openvino/open_model_zoo/models/intel/'")
    args.add_argument("-mc", "--model_acc", type=str, default="FP16", 
                        help="Optional. Model accuracy" "Default value is 'FP16'")
    args.add_argument("-c", "--camera", default='Pi', type=str, 
                        help="Optional. Type of the camera, 'Pi' or 'USB'." "Default value is ''Pi''")
    args.add_argument("-id", "--infer_delay", default=1, type=float, 
                        help="Optional. The interval time (second) between each inference." "Default value is 0.1")
    args.add_argument("-fd", "--find_delay", default=600, type=int, 
                        help="Optional. The interval time (second) between finding & storing & sending somebody. Increase if reached for multiple times." "Default value is 600")
    # args.add_argument("-fddt", "--find_delay_inc_time", default=3, type=int, 
    #                     help="Optional. The number of times of reaching find_delay before doubling find_delay." "Default value is 5")
    args.add_argument("-od", "--output_dir", type=str,
                        default="/home/pi/Pictures/supervise/",
                        help="Optional. Path to captured image directory."
                        "Default value is '/home/pi/Pictures/supervise/'")
    args.add_argument("-ci", "--company_id", type=str,
                        default="ww1d726df782afe2ef",
                        help="Optional. Company id of wechat."
                        "Default value is 'ww1d726df782afe2ef'")
    args.add_argument("-s", "--secret", type=str,
                        default="zxGEelm2ZaC3s04FScAg8tb8235plrkF-aJ5No9eLT0",
                        help="Optional. Secret of wechat."
                        "Default value is 'zxGEelm2ZaC3s04FScAg8tb8235plrkF-aJ5No9eLT0'")

    return parser.parse_args()



def main():
    # ----------------- configure loger ----------------- #
    # check log folder
    log_dir = ".log"
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    # configure logger for write file & print
    logger = myLog("{}/SuperviseLog".format(log_dir))

    args = parse_args()

    # ----------------- prepare openvino model ----------------- #
    # (.xml and .bin files) or (.onnx file)
    model_path = "{}/{}/{}/{}.xml".format(args.model_dir, args.model_name, args.model_acc, args.model_name)

    # initiate model
    openvino_model  = OpenvinoModel(model_path)
    # read model from pre-trained model
    read_fb = openvino_model.readModel()
    if read_fb:
        logger.info("Reading the network: {}".format(args.model_name))
    else:
        logger.error("Read network failed! Check model file.")
        
    
    # configure model
    openvino_model.configureModel()
    # load model into NCS2
    openvino_model.loadModel()


    # ----------------- prepare message API ----------------- #
    message = Message(args.secret, args.company_id)


    old_date = None
    last_cap_time = None
    capture_count = 0
    warning_count = 0
    find_delay_count = 0
    find_delay = args.find_delay

    if args.camera == "USB":

        # ----------------- prepare USB camera ----------------- #
        # USB camera
        cap = cv2.VideoCapture(0)
        logger.info("[info] W, H, FPS: {}, {}, {}".format(cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT), cap.get(cv2.CAP_PROP_FPS)))

        # ----------------- start to loop ----------------- #
        # old_date = None
        # last_cap_time = None
        # capture_count = 0
        # warning_count = 0
        # find_delay_count = 0
        # find_delay = args.find_delay

        # for USB camera
        while(True):
            # check Pi status
            log_message = checkPi()
            if log_message is not None:
                logger.warning(log_message)
                warning_count += 1
                if warning_count >= 3:
                    exit()

            # tik for debug
            start_time = time.time()

            if old_date is None:
                # get current datetime
                old_date = datetime.datetime.now()
            cur_date = datetime.datetime.now()
            if cur_date.date() != old_date.date():
                capture_count = 0
                # reinitiate find_delay every day
                find_delay = args.find_delay

            # read image, USB camera
            ret, image = cap.read()
            if ret is False:
                logger.error("Can not capture image! Check camera.")
                return

            # save for output
            output_image = image.copy()
            # pre process image
            image_process = openvino_model.processInput(image)
            # inference
            detection = openvino_model.inference(image_process)
            # process detection result
            _, _, confidence, xmin, ymin, xmax, ymax = detection
            # get original size
            h, w, _ = output_image.shape

            # check whether there is some one
            if confidence > 0.5:
                now_cap_time = time.time()
                # avoid frequent inference in a short time
                if last_cap_time is None:
                    last_cap_time = time.time()
                # not finding sombody within an interval
                elif now_cap_time - last_cap_time < find_delay:
                    # find_delay_count += 1

                    # inference interval
                    time.sleep(args.infer_delay)
                    # tik for debug
                    print("{:.3f} s".format(time.time() - start_time))
                    print("skip finding, times: {}".format(find_delay_count))

                    # if find_delay_count >= args.find_delay_inc_time:
                    #     find_delay *= 2
                    continue
                last_cap_time = now_cap_time
                # find somebody
                logger.info("Find somebody!")
                capture_count += 1

                # locate bounding box
                xmin = int(xmin * w)
                ymin = int(ymin * h)
                xmax = int(xmax * w)
                ymax = int(ymax * h)

                # add bounding box
                cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                # store image
                output_filename = args.output_dir + "out_{}_{}.png".format(cur_date.date(), capture_count)
                cv2.imwrite(output_filename, output_image)

                if os.path.exists(output_filename):
                    logger.info("Image {} created!".format(output_filename))
                    # send pimage through wechat API
                    send_fb = message.sendMsg(output_filename)
                    if send_fb:
                        logger.info(">> Send => ok")
                    else:
                        logger.error("xx Send => failure! Check API.")
                        return
                else:
                    logger.error("Image {} was not created! Check your permissions.".format(output_filename))
                    return

            # inference interval
            time.sleep(args.infer_delay)

            # tik for debug
            print("{:.3f} s".format(time.time() - start_time))
    else:
        # ----------------- prepare Pi camera ----------------- #
        # Pi camera
        with PiCamera() as cap:
            with PiRGBArray(cap) as raw_cap:
                # old_date = None
                # last_cap_time = None
                # capture_count = 0
                # warning_count = 0
                # find_delay_count = 0
                # find_delay = args.find_delay

                # ----------------- start to loop ----------------- #

                # for Pi camera
                for frame in cap.capture_continuous(raw_cap, format="bgr"):
                    # check Pi status
                    log_message = checkPi()
                    if log_message is not None:
                        logger.warning(log_message)
                        warning_count += 1
                        if warning_count >= 3:
                            exit()

                    # tik for debug
                    start_time = time.time()

                    if old_date is None:
                        # get current datetime
                        old_date = datetime.datetime.now()
                    cur_date = datetime.datetime.now()
                    if cur_date.date() != old_date.date():
                        capture_count = 0
                        # reinitiate find_delay every day
                        find_delay = args.find_delay

                    # read image, Pi camera
                    image = frame.array

                    # save for output
                    output_image = image.copy()
                    # pre process image
                    image_process = openvino_model.processInput(image)
                    # inference
                    detection = openvino_model.inference(image_process)
                    # process detection result
                    _, _, confidence, xmin, ymin, xmax, ymax = detection
                    # get original size
                    h, w, _ = output_image.shape

                    # check whether there is some one
                    if confidence > 0.5:
                        now_cap_time = time.time()
                        # avoid frequent inference in a short time
                        if last_cap_time is None:
                            last_cap_time = time.time()
                        # not finding sombody within an interval
                        elif now_cap_time - last_cap_time < find_delay:
                            
                            find_delay_count += 1

                            # need to flush memery of camera
                            raw_cap.truncate(0)
                            print("skip find, times: {}".format(find_delay_count))

                            # if find_delay_count >= args.find_delay_inc_time:
                            #     find_delay_count = 0
                            #     find_delay += args.find_delay
                            #     print("increase find delay, now: {}".format(find_delay))
                                
                            # inference interval
                            time.sleep(args.infer_delay)
                            # tik for debug
                            print("{:.3f} s".format(time.time() - start_time))
                            continue

                        # when not keep skipping find someone, reinitiate find_delay_count
                        find_delay_count = 0
                        last_cap_time = now_cap_time
                        # find somebody
                        logger.info("Find somebody!")
                        capture_count += 1

                        # locate bounding box
                        xmin = int(xmin * w)
                        ymin = int(ymin * h)
                        xmax = int(xmax * w)
                        ymax = int(ymax * h)

                        # add bounding box
                        cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                        # store image
                        output_filename = args.output_dir + "out_{}_{}.png".format(cur_date.date(), capture_count)
                        cv2.imwrite(output_filename, output_image)

                        if os.path.exists(output_filename):
                            logger.info("Image {} created!".format(output_filename))
                            # send image through wechat API
                            send_fb = message.sendMsg(output_filename)
                            if send_fb:
                                logger.info(">> Send => ok")
                            else:
                                logger.error("xx Send => failure! Check API.")
                                return
                        else:
                            logger.error("Image {} was not created! Check your permissions.".format(output_filename))
                            return

                    # need to flush memery of camera
                    raw_cap.truncate(0)
                    # raw_cap.seek(0)

                    # inference interval
                    time.sleep(args.infer_delay)


                    # tik for debug
                    print("{:.3f} s".format(time.time() - start_time))


if __name__ == "__main__":
    main()



